import os
import random
import glob
from PIL import Image, ImageEnhance, ImageOps

DESKTOP_PATH = os.path.join(os.path.expanduser("~"), "Desktop")
OUTPUT_DIR = os.path.join(DESKTOP_PATH, "CNN_model")

# 配置参数
NORMAL_IMAGES_DIR = os.path.join(OUTPUT_DIR,"phone_screen_dataset", "normal")  # 原始正常图片目录
SCREEN_ART_DIR = os.path.join(OUTPUT_DIR, "glitch_patterns")  # 花屏图片目录
ABNORMAL_DIR = os.path.join(OUTPUT_DIR, "CNN_test", "abnormal")  # 异常图片目录
mask_dir = os.path.join(OUTPUT_DIR, "phone_screen_dataset", "mask")  # 异常图片掩码目录
IMAGES_PER_ORIGINAL = 20  # 每张原始图片生成的增强图片数量
#ART_PER_IMAGE = 1  # 每张图片叠加的花屏数量

os.makedirs(ABNORMAL_DIR, exist_ok=True)

def load_images_from_dir(directory, extensions=('png','jpg')):
    images = []
    for ext in extensions:
        for img_path in glob.glob(os.path.join(directory, f'*.{ext}')):
            try:
                img = Image.open(img_path).convert("RGBA")
                images.append((img, os.path.basename(img_path)))
                print(f"加载图片: {img_path}")
            except Exception as e:
                print(f"加载图片失败 {img_path}: {e}")
    return images

def overlay_single_art(normal_img, art_img, generate_mask=True):
    """叠加单张花屏，并同步生成掩码"""
    img = normal_img.copy()
    h, w = img.size[1], img.size[0]  # 高度、宽度
    
    # 调整大小
    art_resized = art_img.resize((w, h), Image.Resampling.LANCZOS)

    #调整透明度
    alpha_factor = 1.0
    art_with_alpha = art_resized.copy()
    alpha = art_with_alpha.split()[3]
    #alpha = ImageEnhance.Brightness(alpha).enhance(alpha_factor)
    art_with_alpha.putalpha(alpha)

    #旋转
    rotate_angle = 0
    art_rotated = art_with_alpha.rotate(rotate_angle, expand=True, fillcolor=(0, 0, 0, 0))

    #调整位置
    max_x = img.width - art_rotated.width
    max_y = img.height - art_rotated.height
    x = random.randint(0, max_x) if max_x > 0 else 0
    y = random.randint(0, max_y) if max_y > 0 else 0
    
    # 叠加花屏到原图
    img.paste(art_rotated, (x, y), art_rotated)
    
    # 生成掩码：花屏区域为255，背景为0
    if generate_mask:
        mask = Image.new('L', (w, h), 0)  # 单通道黑色图像
        # 提取花屏的alpha通道（非透明区域即为异常）
        art_alpha = art_rotated.split()[3]  # alpha通道（0为透明，255为不透明）
        # 将花屏的非透明区域绘制到掩码上
        mask.paste(255, (x, y), art_alpha)  # 花屏区域设为255
        return img, mask
    else:
        return img

def generate_abnormal_images(normal_images, art_images, mask_dir="masks"):
    os.makedirs(mask_dir, exist_ok=True)  # 创建掩码文件夹
    
    for img, img_name in normal_images:
        selected_arts = random.sample(art_images, 1)  # 每次随机抽选1种花屏
        base_name = os.path.splitext(img_name)[0]
        for i, (art_img, art_name) in enumerate(selected_arts):
            # 叠加花屏并生成掩码
            abnormal_img, mask = overlay_single_art(img, art_img, generate_mask=True)
            
            # 保存异常图像
            img_path = os.path.join(ABNORMAL_DIR, f"{base_name}_abnormal_{i}.png")
            abnormal_img.convert("RGB").save(img_path)
            
            # 保存对应掩码
            mask_path = os.path.join(mask_dir, f"{base_name}_abnormal_{i}_mask.png")
            mask.save(mask_path)  # 掩码为单通道灰度图
            print(f"生成异常图像和掩码: {img_path}, {mask_path}")

def main():
    # 加载原始图片和花屏图片
    print("加载原始正常图片...")
    normal_images = load_images_from_dir(NORMAL_IMAGES_DIR)
    
    if not normal_images:
        print("错误: 没有找到原始正常图片，请检查NORMAL_IMAGES_DIR目录")
        return
    
    #print("加载花屏图片...")
    screen_art_images = load_images_from_dir(SCREEN_ART_DIR)
    
    # 生成异常图片
    print("开始生成异常图片...")
    generate_abnormal_images(normal_images, screen_art_images)
    print(f"已生成 {len(normal_images) * 1} 张异常图片")
    
    print("数据集生成完成!")

if __name__ == "__main__":
    main()
